import matplotlib.pyplot as plt
import numpy as np
import basic_func.general_func as bf

def zh_font(params):
    if params['zh_font'] == 'yes':
        plt.rcParams['font.sans-serif']=['SimHei']
        plt.rcParams['axes.unicode_minus']=False

def whether_label(color,ax,params,rects):
    if params['bar_label']=='yes':
        r,g,b,_=color
        # r*g*b 用来估计颜色的亮度
        text_color = 'white' if r * g * b < 0.5 \
                     else 'darkgrey'
        fmt = set_label_text_format(params)
        ax.bar_label(rects, label_type='center',
                     color=text_color,
                     fmt = fmt)
def get_color(params,data):
    color_name = params['cmap_color']
    colors = plt.colormaps[color_name]
    colors = colors(np.linspace(0,1,data.shape[1]))
    return colors

def set_label_text_format(params):
    if params['label_format']==1:
        fmt = '%d%%'
    elif params['label_format']==2:
        fmt = '%g'
    return fmt

def survey_bar(df,params):
    zh_font(params)
    data = df.values
    cum_data = data.cumsum(axis =1)
    category_colors = get_color(params,data)
    fig,ax = plt.subplots(figsize=(10,6))
    # 使用numpy 数组向量化计算
    for i, colname in enumerate(df.columns):
        rects = ax.barh(range(len(df)),width = data[:,i],
                        left = cum_data[:,i]-data[:,i],
                        height = 0.6,
                        color = category_colors[i])
        whether_label(category_colors[i],ax,params,rects)
        
    ax.xaxis.set_visible(False)
    ax.set_yticks(range(len(df)),labels = df.index)
    ax.legend(df.columns,ncols = len(df.columns),
              bbox_to_anchor = (0,1),
              loc = 'lower left',
              fontsize = 13)
    plt.show()


if __name__=='__main__':
    df,params = bf.read_data('test_data/水平条形图展示调查问卷.xlsx')
    survey_bar(df,params)
